# -*- coding:utf-8 -*-

import json
import config
from loguru import logger
from kafka import KafkaProducer


class Producer():

    def __init__(self, topic, bootstrap_servers=config.SERVERS):
        """
        kafka 生产者
        :param bootstrap_servers: 地址
        """
        self.producer = KafkaProducer(
            bootstrap_servers=bootstrap_servers,
            value_serializer=lambda m: json.dumps(m).encode(),
            acks='all')  # json 格式化发送的内容
        self.topic = topic

    def sync_producer(self, data_li: list):
        """
        同步发送 数据
        :param data_li:  发送数据
        :return:
        """
        for data in data_li:
            future = self.producer.send(self.topic, data)
            record_metadata = future.get(timeout=10)  # 同步确认消费
            partition = record_metadata.partition  # 数据所在的分区
            offset = record_metadata.offset  # 数据所在分区的位置
            print('save success, partition: {}, offset: {}'.format(
                partition, offset))

    def asyn_producer(self, data_li: list):
        """
        异步发送数据
        :param data_li:发送数据
        :return:
        """
        for data in data_li:
            self.producer.send(self.topic, data)
        self.producer.flush()  # 批量提交

    def asyn_producer_callback(self, data_li: list, key: str = None):
        """
        异步发送数据 + 发送状态处理
        :param data_li:发送数据
               key: 发送数据的key
        :return:
        """
        for data in data_li:
            self.producer.send(self.topic, value=data, key=key).add_callback(
                self.send_success).add_errback(self.send_error)
        self.producer.flush()  # 批量提交

    def send_success(self, *args, **kwargs):
        """异步发送成功回调函数"""
        logger.info('save success')
        return

    def send_error(self, *args, **kwargs):
        """异步发送错误回调函数"""
        logger.error('save error')
        return

    def close_producer(self):
        try:
            self.producer.close()
        except Exception:
            pass


if __name__ == '__main__':
    topic = "test"
    kp = Producer(topic)
    try:
        send_data_li = [
            config.KafkaMsg(id=x, detail=12) for x in range(100, 120)
        ]
        kp.asyn_producer_callback(send_data_li)
    except Exception as e:
        logger.exception(f'except error {e}')
    finally:
        kp.close_producer()
